from structAIDetect.modeltrain.disinfermodel import DisplaceInterpolator
from structAIDetect.modeltrain.structinfermodel import StructReconstruct
from structAIDetect.modeltrain.modaldataset import ModalDataset
import os
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import matplotlib.pyplot as plt
import argparse


LR = 0.001
BATCH_SIZE = 16
INPUT_LEN = 44
HIDDEN_SIZE = 128
NUMLAYERS = 1
TOTAL_MODAL_NUM = 4
EPOCHS = 2000


def displace_predict(displace_model, modalinfo_edge, pred_len):
    displace_model.eval()
    pred_result = torch.zeros([modalinfo_edge.shape[0], pred_len, modalinfo_edge.shape[2] - TOTAL_MODAL_NUM])
    h_t = torch.zeros(NUMLAYERS, modalinfo_edge.shape[0], HIDDEN_SIZE)
    c_t = torch.zeros(NUMLAYERS, modalinfo_edge.shape[0], HIDDEN_SIZE)
    next_input = modalinfo_edge[:, [0], TOTAL_MODAL_NUM:]
    for pred_idx in range(pred_len):
        next_input = torch.concat((modalinfo_edge[:, [0], :TOTAL_MODAL_NUM], next_input), axis=2)
        outputs, (h_t, c_t) = displace_model(next_input, (h_t, c_t))
        next_input = outputs.unsqueeze(1)
        pred_result[:, [pred_idx], :] = outputs.unsqueeze(1)
    
    return pred_result, h_t


def predict_loss(model, displace_model, testloader, criterion, pred_len=10):
    model.eval()
    losses = 0
    with torch.no_grad():
        for modalinfo, targets in testloader:
            modalinfo_edge = modalinfo[:, [0], :]
            displacement, displace_hidden = displace_predict(displace_model, modalinfo_edge, pred_len)
            cellstate_predict = model(displace_hidden[0, :, :], modalinfo[:, 0, :4])
            cellstate_predict = cellstate_predict[:, :, :41, :]
            targets = targets.squeeze(3)
            loss = criterion(cellstate_predict, targets)
            losses += loss

        return losses


def train_model(model, displace_model, trainloader, criterion, optimizer, pred_len=10):
    model.train()
    losses = 0
    for modalinfo, targets in trainloader:
        optimizer.zero_grad()
        modalinfo_edge = modalinfo[:, [0], :]
        displacement, displace_hidden = displace_predict(displace_model, modalinfo_edge, pred_len)
        cellstate_predict = model(displace_hidden[0, :, :], modalinfo[:, 0, :4])
        cellstate_predict = cellstate_predict[:, :, :41, :]
        targets = targets.squeeze(3)
        loss = criterion(cellstate_predict, targets)
        loss.backward()
        losses += loss
        optimizer.step()

    return losses


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--epochs', type=int, default=2000, help='输入训练轮次')
    args = parser.parse_args()
    EPOCHS = args.epochs
    
    displacement_model = DisplaceInterpolator(input_size=INPUT_LEN, hidden_size=HIDDEN_SIZE, num_layers=NUMLAYERS)
    displacement_model.load_state_dict(torch.load(os.path.join('model', 'displace_predict_para.pth')))

    datapath = os.path.join(os.getcwd(), "data", "norm")

    model = StructReconstruct(HIDDEN_SIZE, 4)
    criterion = nn.CrossEntropyLoss(weight=torch.tensor([0.8, 0.2]))
    optimizer = torch.optim.Adam(model.parameters(), lr=LR)

    train_data = ModalDataset(datapath, train_test='train')
    test_data = ModalDataset(datapath, train_test='test')
    train_loader = DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True)
    test_loader = DataLoader(test_data, batch_size=BATCH_SIZE, shuffle=True)

    train_losses = []
    test_losses = []
    for epoch in range(EPOCHS):
        train_loss = train_model(model, 
                                displacement_model, 
                                train_loader, 
                                criterion, 
                                optimizer)
        train_losses.append(train_loss.item())
        test_loss = predict_loss(model, 
                                displacement_model, 
                                test_loader, 
                                criterion)
        test_losses.append(test_loss.item())
        print(f"Epoch {epoch+1}, Train Loss: {train_loss:.4f}")
        print(f"Epoch {epoch+1}, Test Loss: {test_loss:.4f}")

    plt.figure()
    plt.plot(train_losses)
    plt.title("Train Loss")
    plt.show()
    plt.figure()
    plt.plot(test_losses)
    plt.title("Test Loss")
    plt.show()

    savepath = os.path.join(os.getcwd(), 'model')
    torch.save(model.state_dict(), os.path.join(savepath, 'struct_reconstruct_para.pth'))

if __name__ == '__main__':
    main()